import base64
import json
import os

import requests
from typing import Optional, Dict, Any


class ElasticSearchClient:
    def __init__(self, config_file: str = 'config.json'):
        # 从配置文件加载 Elasticsearch 的配置
        config_path = os.path.join(os.path.dirname(__file__), config_file)
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        self.elasticsearch_url = config.get("elasticsearch_url")
        self.elasticsearch_login = config.get("elasticsearch_login")
        # 设置超时时间（秒）
        self.timeout = 30

    def _get_auth_header(self) -> Dict[str, str]:
        """生成认证头"""
        auth_bytes = self.elasticsearch_login.encode('utf-8')
        auth_b64 = base64.b64encode(auth_bytes).decode('utf-8')
        return {
            'Authorization': f'Basic {auth_b64}',
            'Content-Type': 'application/json'
        }


    def execute_url1(
            self,
            path: str,
            method: str = 'GET',
            body: Optional[Dict] = None,
            params: Optional[Dict] = None
    ) -> Optional[Dict[str, Any]]:
        """执行ES查询"""
        try:
            url = f"{self.elasticsearch_url}{path}"
            headers = {'Content-Type': 'application/json'}

            response = requests.request(
                method=method,
                url=url,
                headers=headers,
                data=json.dumps(body) if body else None,
                params=params,
                timeout=self.timeout
            )

            if response.status_code != 200:
                print(f"ES error response: {response.text}")
                return None

            return response.json()

        except Exception as e:
            print(f"ES request error: {str(e)}")
            return None


    def execute_url(self, url_path: str, method: str = 'GET', body: Optional[Dict] = None) -> Optional[Dict[str, Any]]:
        """
        执行 ElasticSearch 请求

        Args:
            url_path: API路径
            method: 请求方法 (GET/POST/PUT/DELETE)
            body: 请求体数据（可选）

        Returns:
            Dict: 响应数据
            None: 如果请求失败
        """
        try:
            full_url = f"{self.elasticsearch_url}{url_path}"
            headers = self._get_auth_header()

            response = requests.request(
                method=method,
                url=full_url,
                headers=headers,
                json=body,
                timeout=self.timeout
            )

            response.raise_for_status()  # 检查响应状态
            return response.json()

        except requests.exceptions.RequestException as e:
            print(f"Error executing ElasticSearch request: {str(e)}")
            return None


    def search(self, index: str, query: Dict) -> Optional[Dict[str, Any]]:
        """
        执行搜索查询

        Args:
            index: 索引名称
            query: 查询条件

        Returns:
            Dict: 查询结果
            None: 如果查询失败
        """
        url_path = f"/{index}/_search"
        return self.execute_url(url_path, method='GET', body=query)

    def get_alarm_statistics(self, alarm_grades: str) -> Dict[str, Any]:
        """获取告警统计数据"""
        grade_list = [int(grade) for grade in alarm_grades.split(',')] if alarm_grades else None

        query = {
            "size": 0,
            "aggs": {
                "alarm_grade_stats": {
                    "terms": {
                        "field": "alarm_grade",
                        "size": 10
                    }
                }
            }
        }

        if grade_list:
            query["query"] = {
                "bool": {
                    "must": [
                        {"terms": {"alarm_grade": grade_list}}
                    ]
                }
            }

        return self.execute_url("/gd-zhiheyunv2-alarm/_search", "POST", body=query)

    def get_alarm_list(self, query_params: Dict) -> Dict[str, Any]:
        """获取告警列表数据"""
        # 使用与find_alarms_data相同的查询逻辑
        return self.execute_url("/gd-zhiheyunv2-alarm/_search", "POST", body=query_params)
